"""
统一客户端类
"""

from rich.console import Console
from rich.prompt import Prompt
from rich.panel import Panel
from rich.table import Table
from agent_mode import AgentMode
from control_mode import ControlMode
from relay_api import RelayAPI

console = Console()

class SSHRemoteClient:
    def __init__(self):
        self.agent_mode = AgentMode()
        self.control_mode = ControlMode()
        self.api = RelayAPI()
    
    def run(self):
        """运行客户端主程序"""
        self._show_banner()
        
        while True:
            choice = self._show_main_menu()
            
            if choice == "1":
                self._run_agent_mode()
            elif choice == "2":
                self._run_control_mode()
            elif choice == "3":
                self._show_status()
            elif choice == "4":
                console.print("[bold green]感谢使用SSH远程控制工具！[/bold green]")
                break
            else:
                console.print("[bold red]无效选择，请重新输入[/bold red]")
    
    def _show_banner(self):
        """显示程序横幅"""
        banner = """
[bold cyan]
╔═══════════════════════════════════════╗
║        SSH 远程控制工具 v1.0          ║
║     支持内网穿透的安全远程连接        ║
╚═══════════════════════════════════════╝
[/bold cyan]
        """.strip()
        
        console.print(banner)
        console.print()
    
    def _show_main_menu(self) -> str:
        """显示主菜单"""
        menu_text = """
[bold white]请选择工作模式:[/bold white]

[bold green]1.[/bold green] 被控模式 - 等待他人连接
[bold blue]2.[/bold blue] 控制模式 - 连接他人
[bold yellow]3.[/bold yellow] 查看服务器状态
[bold red]4.[/bold red] 退出程序
        """.strip()
        
        panel = Panel(
            menu_text,
            title="[bold cyan]主菜单[/bold cyan]",
            border_style="cyan",
            padding=(1, 2)
        )
        
        console.print(panel)
        
        choice = Prompt.ask(
            "[bold cyan]请输入选项 (1-4)[/bold cyan]",
            choices=["1", "2", "3", "4"],
            default="1"
        )
        
        return choice
    
    def _run_agent_mode(self):
        """运行被控模式"""
        try:
            self.agent_mode.start()
        except Exception as e:
            console.print(f"[bold red]被控模式运行错误: {str(e)}[/bold red]")
        
        console.print("\n按 [bold yellow]Enter[/bold yellow] 返回主菜单...")
        input()
    
    def _run_control_mode(self):
        """运行控制模式"""
        try:
            self.control_mode.start()
        except Exception as e:
            console.print(f"[bold red]控制模式运行错误: {str(e)}[/bold red]")
        
        console.print("\n按 [bold yellow]Enter[/bold yellow] 返回主菜单...")
        input()
    
    def _show_status(self):
        """显示服务器状态"""
        console.print("\n[bold cyan]正在获取服务器状态...[/bold cyan]")
        
        result = self.api.get_status()
        
        if result.get("status") == "running":
            # 创建状态表格
            table = Table(title="中继服务器状态", show_header=True, header_style="bold magenta")
            table.add_column("项目", style="cyan", no_wrap=True)
            table.add_column("值", style="green")
            
            table.add_row("服务器状态", "运行中 ✓")
            table.add_row("总连接数", str(result.get("total_connections", 0)))
            table.add_row("活跃连接数", str(result.get("active_connections", 0)))
            table.add_row("服务器时间", result.get("server_time", "未知"))
            
            console.print(table)
        else:
            console.print(f"[bold red]获取状态失败: {result.get('message', '未知错误')}[/bold red]")
        
        console.print("\n按 [bold yellow]Enter[/bold yellow] 返回主菜单...")
        input()

